package com.dragon.executor;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.lang.management.ManagementFactory;
import java.lang.management.MemoryMXBean;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 包装器，通过启动一个独立的子进程来执行另一个 main 方法，同时跟踪峰值内存使用情况。
 */
public class MemoryTrackingExecutor {

    private static final String SEPARATOR = "---MEMORY_SEPARATOR---";

    public static void main(String[] args) {
        if (args.length < 1) {
            System.err.println("用法: MemoryTrackingExecutor <真正的main类> [参数...]");
            System.exit(1);
        }

        String realMainClassName = args[0];
        String[] realArgs = Arrays.copyOfRange(args, 1, args.length);

        final MemoryMXBean memoryMXBean = ManagementFactory.getMemoryMXBean();
        final AtomicLong peakMemory = new AtomicLong(0);

        // 内存监控线程池
        ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();

        // Shutdown Hook 负责在最后打印内存
        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            scheduler.shutdownNow();
            long usedHeapMemory = memoryMXBean.getHeapMemoryUsage().getUsed();
            peakMemory.getAndUpdate(x -> Math.max(x, usedHeapMemory));
            System.out.println();
            System.out.println(SEPARATOR);
            System.out.println(peakMemory.get() / 1024);
        }));

        scheduler.scheduleAtFixedRate(() -> {
            try {
                long usedHeapMemory = memoryMXBean.getHeapMemoryUsage().getUsed();
                peakMemory.getAndUpdate(x -> Math.max(x, usedHeapMemory));
            } catch (Exception e) {
                // 忽略
            }
        }, 0, 20, TimeUnit.MILLISECONDS);

        Process subProcess = null;
        try {
            // 创建一个新的子进程来运行用户代码
            List<String> command = new ArrayList<>();
            command.add("java");
            command.add("-Xmx256m");
            command.add("-Dfile.encoding=UTF-8");
            command.add("-cp");
            command.add(System.getProperty("java.class.path"));
            command.addAll(Arrays.asList(args));

            // 使用与当前 Executor 相同的 classpath
            command.add(System.getProperty("java.class.path"));
            command.add(realMainClassName);
            command.addAll(Arrays.asList(realArgs));

            ProcessBuilder processBuilder = new ProcessBuilder(command);
            subProcess = processBuilder.start();

            // 将 Executor 的标准输入 转发给 子进程的标准输入
            Thread inputForwarder = new Thread(new StreamForwarder(System.in, subProcess.getOutputStream()));
            // 将 子进程的标准输出 转发给 Executor 的标准输出
            Thread outputForwarder = new Thread(new StreamForwarder(subProcess.getInputStream(), System.out));
            // 将 子进程的错误输出 转发给 Executor 的错误输出
            Thread errorForwarder = new Thread(new StreamForwarder(subProcess.getErrorStream(), System.err));

            inputForwarder.start();
            outputForwarder.start();
            errorForwarder.start();

            // 等待子进程执行完毕
            int exitCode = subProcess.waitFor();

            // 等待 I/O 转发线程结束
            inputForwarder.join();
            outputForwarder.join();
            errorForwarder.join();

            // 退出时使用子进程的退出码
            System.exit(exitCode);

        } catch (Exception e) {
            e.printStackTrace();
            System.exit(1);
        } finally {
            if (subProcess != null) {
                subProcess.destroyForcibly();
            }
        }
    }

    /**
     * 一个简单的工具类，用于将一个输入流的内容持续地复制到另一个输出流。
     */
    static class StreamForwarder implements Runnable {
        private final InputStream in;
        private final OutputStream out;

        StreamForwarder(InputStream in, OutputStream out) {
            this.in = in;
            this.out = out;
        }

        @Override
        public void run() {
            try {
                byte[] buffer = new byte[1024];
                int len;
                while ((len = in.read(buffer)) != -1) {
                    out.write(buffer, 0, len);
                    out.flush();
                }
            } catch (IOException e) {
                // 忽略管道关闭等异常
            }
        }
    }
}